package cc.git.liuyan;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;
import java.security.ProtectionDomain;
import java.util.ArrayList;
import java.util.List;

public class Agent {
    private static List<String> notNeedTransFormClassNamePrefix = new ArrayList<String>();
    private static boolean startTransForm = false;
    private static List<String> needTranformPrefex = new ArrayList<>();
    private static boolean saveModifiedBytecode = true;
    private static File modifiedBytecodeSaveDir = null;

    static {
        notNeedTransFormClassNamePrefix.add("cc/git/liuyan/Agent");
        notNeedTransFormClassNamePrefix.add("cc/git/liuyan/AgentMain");
        notNeedTransFormClassNamePrefix.add("cc/git/liuyan/SelfLog");
        notNeedTransFormClassNamePrefix.add("cc/git/liuyan/SelfOutPutStream");
        notNeedTransFormClassNamePrefix.add("sun");
        notNeedTransFormClassNamePrefix.add("com/sun");
        notNeedTransFormClassNamePrefix.add("javax");
        notNeedTransFormClassNamePrefix.add("java");
        notNeedTransFormClassNamePrefix.add("null");
        notNeedTransFormClassNamePrefix.add("jdk");
        notNeedTransFormClassNamePrefix.add("org/omg");
        notNeedTransFormClassNamePrefix.add("com/intellij");
        notNeedTransFormClassNamePrefix.add("com/zeroturnaround");
        notNeedTransFormClassNamePrefix.add("org/springframework");
        notNeedTransFormClassNamePrefix.add("org/slf4j");
        notNeedTransFormClassNamePrefix.add("com/mysql");
        notNeedTransFormClassNamePrefix.add("com/alibaba");
        notNeedTransFormClassNamePrefix.add("com/baomidou");
        notNeedTransFormClassNamePrefix.add("ch/qos");
        notNeedTransFormClassNamePrefix.add("com/google");
        notNeedTransFormClassNamePrefix.add("com/netflix");
        notNeedTransFormClassNamePrefix.add("com/jayway");
        notNeedTransFormClassNamePrefix.add("com/github");
        notNeedTransFormClassNamePrefix.add("com/fasterxml");
        notNeedTransFormClassNamePrefix.add("com/aliyun");
        notNeedTransFormClassNamePrefix.add("com/itextpdf");
        notNeedTransFormClassNamePrefix.add("com/qcloud");
        notNeedTransFormClassNamePrefix.add("com/zaxxer");
        notNeedTransFormClassNamePrefix.add("com/typesafe");

        if (saveModifiedBytecode) {
            modifiedBytecodeSaveDir = new File("./localUseAgentByteCode");
            if (modifiedBytecodeSaveDir.exists() && modifiedBytecodeSaveDir.isDirectory()) {
                File[] files = modifiedBytecodeSaveDir.listFiles();
                if (files != null) {
                    for (File file : files) {
                        file.delete();
                    }
                }
                modifiedBytecodeSaveDir.delete();
            }
        }
    }

    public static void premain(String agentArgs, Instrumentation inst) {
        SelfLog.info("Agent premain method called with args: " + agentArgs);
        inst.addTransformer(new MethodPrintExcuteTimeTransformer(), true);
    }

    public static void agentmain(String agentArgs, Instrumentation inst) {
        SelfLog.info("Agent agentmain method called with args -> main class: " + agentArgs);
        String longName = agentArgs.substring(0, agentArgs.lastIndexOf("."));
        longName = longName.replace(".", "/");
        needTranformPrefex.add(longName);
        while (longName.contains("/")) {
            longName = longName.substring(0, longName.lastIndexOf("/"));
            needTranformPrefex.add(longName);
        }
        startTransForm = true;
        SelfLog.info("agentmain needTranformPrefex: " + String.join(", ", needTranformPrefex));
        SelfLog.info("isRetransformClassesSupported：" + inst.isRetransformClassesSupported());
        SelfLog.info("isRedefineClassesSupported：" + inst.isRedefineClassesSupported());
        inst.addTransformer(new MethodPrintExcuteTimeTransformer(), true);

        Class classes[] = inst.getAllLoadedClasses();
        for (int i = 0; i < classes.length; i++) {
            if (inst.isModifiableClass(classes[i])) {
                try {
                    inst.retransformClasses(classes[i]);
                } catch (UnmodifiableClassException e) {
                    SelfLog.info("Cannot retransform class: " + classes[i].getName());
                }
            }
        }
    }

    public static class MethodPrintExcuteTimeTransformer implements ClassFileTransformer {
        public MethodPrintExcuteTimeTransformer() {
            SelfLog.info("MethodPrintExcuteTimeTransformer created");
        }

        @Override
        public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
            try {
                ClassPool pool = ClassPool.getDefault();
                CtClass cc = pool.makeClass(new java.io.ByteArrayInputStream(classfileBuffer));
                CtMethod[] methods = cc.getDeclaredMethods();
                if (className == null || methods == null || methods.length == 0 || cc.isInterface() || cc.isEnum() || cc.isAnnotation()) {
                    return classfileBuffer;
                }
                if (classNameNotNeedTransForm(className)) {
                    methods = null;
                    cc = null;
                    pool = null;
                    return classfileBuffer;
                } else if (!startTransForm) {
                    for (CtMethod method : methods) {
                        String longName = method.getLongName();
                        if (longName.endsWith(".main(java.lang.String[])")) {
                            longName = longName.replace(".main(java.lang.String[])", "");
                            longName = longName.substring(0, longName.lastIndexOf("."));
                            longName = longName.replace(".", "/");
                            needTranformPrefex.add(longName);
                            while (longName.contains("/")) {
                                longName = longName.substring(0, longName.lastIndexOf("/"));
                                needTranformPrefex.add(longName);
                            }
                            startTransForm = true;
                            SelfLog.info("premain needTranformPrefex: " + String.join(", ", needTranformPrefex));
                            break;
                        }
                    }
                    return classfileBuffer;
                } else if (needTranformPrefex.size() > 0 && needTranformPrefex.stream().anyMatch(className::startsWith)) {
                    if (className.contains("$")) {
                        //代理类只保存字节码用于分析，不做字节码修改
                        if (saveModifiedBytecode) {
                            File filePath = new File(modifiedBytecodeSaveDir, className.substring(0, className.lastIndexOf("/")).replace("/", File.separator));
                            filePath.mkdirs();
                            File file = new File(filePath, cc.getSimpleName() + ".class");
                            try (FileOutputStream fos = new FileOutputStream(file)) {
                                fos.write(classfileBuffer);
                                fos.flush();
                            } catch (IOException e) {
                                SelfLog.info("Failed to write modified bytecode to file: " + e.getMessage(), e);
                            }
                            file = null;
                            filePath = null;
                        }
                        return classfileBuffer;
                    }
                    SelfLog.info("MethodPrintExcuteTimeTransformer transforming class: " + className);
                    for (CtMethod method : methods) {
                        if (method.isEmpty()) {
                            continue;
                        }
                        method.addLocalVariable("startTime", CtClass.longType);
                        method.insertBefore(
                                "startTime = System.currentTimeMillis(); \n" +
                                        "cc.git.liuyan.SelfLog.printStart(\"Method: " + method.getLongName() + " started at: \" + startTime);"
                        );
                        method.insertAfter("cc.git.liuyan.SelfLog.printEnd(\"Method: " + method.getLongName() + " ended at: \" + System.currentTimeMillis() + \" 执行耗时: \" + (System.currentTimeMillis() - startTime) + \" ms\");");
                    }
                    byte[] modifiedBytecode = cc.toBytecode();
                    if (saveModifiedBytecode) {
                        File filePath = new File(modifiedBytecodeSaveDir, className.substring(0, className.lastIndexOf("/")).replace("/", File.separator));
                        filePath.mkdirs();
                        File file = new File(filePath, cc.getSimpleName() + ".class");
                        try (FileOutputStream fos = new FileOutputStream(file)) {
                            fos.write(modifiedBytecode);
                            fos.flush();
                        } catch (IOException e) {
                            SelfLog.info("Failed to write modified bytecode to file: " + e.getMessage(), e);
                        }
                        file = null;
                        filePath = null;
                    }
                    return modifiedBytecode;
                }
            } catch (Exception e) {
                SelfLog.info(e.getMessage(), e);
            }

            return classfileBuffer;
        }

        private boolean classNameNotNeedTransForm(String className) {
            if (className == null || notNeedTransFormClassNamePrefix.stream().anyMatch(className::startsWith)) {
                return true;
            }
            return false;
        }
    }
}
